import os
import numpy as np

np.set_printoptions(threshold=np.inf)

def GenerateData(path):
    path_file=os.listdir(path)
    labels,images=[],[]
    for file in path_file:
        labels.append(int(file.split('_')[0]))
        single=[]
        with open(os.path.join(path,file),'r') as f:
            for line in f.readlines():
                for l in range(len(line)-1):
                    single.append(int(line[l]))
            images.append(single)
    return np.array(images),np.array(labels)


def TrainNB(train_x,train_y,test_x,test_y):
    c,p=[],[]
    for i in range(10):
        print("i",i)
        s=np.sum(train_x[train_y == i],0)+1
        print(s.shape)
        print(s.reshape(32,32))
        # break
        p.append(s/np.sum(s))
        c.append(np.sum(test_y==i)/len(test_y==i))
    P=np.zeros((len(test_x),10))
    for i,test in enumerate(test_x):
        for j in range(10):
            P[i,j]=np.sum(test*np.log(p[j]))+np.log(c[j])
    print(P.argmax(1))
    print(test_y)
    print('精确率:{:.2f}%'.format(np.mean(np.equal(P.argmax(1),test_y))*100))

test='testDigits/'
train='trainingDigits/'
train_x,train_y=GenerateData(train)
test_x,test_y=GenerateData(test)
TrainNB(train_x,train_y,test_x,test_y)
